package org.heart.thirdService.AIServiceAPI.config;

import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import lombok.extern.slf4j.Slf4j;
import org.heart.thirdService.AIServiceAPI.service.AisServices;
import org.heart.thirdService.AIServiceAPI.tools.WeatherTool;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;


import javax.annotation.Resource;

import java.util.List;

import static dev.langchain4j.model.openai.OpenAiChatModelName.*;

@Slf4j
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    private final UiUiConfig uiUiConfig;

    @Resource
    private WeatherTool weatherTool;

    @Autowired
    public ChatWebSocketHandler(UiUiConfig uiUiConfig) {
        this.uiUiConfig = uiUiConfig;
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {

        log.info("需要发送消息体：{}",message.getPayload());
        log.info("注入的apiKey={}，baseUrl={}",uiUiConfig.getApiKey(),uiUiConfig.getBaseUrl());


        String userMessage = message.getPayload();

        StreamingChatModel model = OpenAiStreamingChatModel.builder()
                /**
                 * API key
                 */
                .apiKey(uiUiConfig.getApiKey())
                /**
                 * GPT 模型
                 */
                .modelName(GPT_4_O)
                .baseUrl(uiUiConfig.getBaseUrl())
                .strictTools(true) // ✅ 添加工具支持（支持多个工具）
                .timeout(java.time.Duration.ofSeconds(5)) // 增加超时时间
                .build();

        AisServices assistant = AiServices.builder(AisServices.class)
                .streamingChatModel(model)
                .tools(List.of(weatherTool))
                .build();


        TokenStream tokenStream = assistant.chat11(userMessage);

        // 流式返回时，实时发送给前端
        tokenStream
                .onPartialResponse(partialResponse -> {
                    try {
                        session.sendMessage(new TextMessage(partialResponse));
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                })
                .onCompleteResponse(response -> {
                    try {
                        // 你可以选择发送最终完整消息或结束信号
                        session.sendMessage(new TextMessage("[DONE]"));
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                })
                .onError(error -> {
                    try {
                        session.sendMessage(new TextMessage("[ERROR]" + error.getMessage()));
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                })
                .start();
//        model.chat(userMessage, new StreamingChatResponseHandler() {
//            @Override
//            public void onPartialResponse(String partialResponse) {
//                try {
//                    session.sendMessage(new TextMessage(partialResponse));
//                } catch (IOException e) {
//                    e.printStackTrace();
//                }
//            }
//
//            @Override
//            public void onCompleteResponse(ChatResponse completeResponse) {
//                try {
//                    session.sendMessage(new TextMessage("[DONE]"));
////                    session.close(); // 关闭连接（可选）
//                } catch (IOException e) {
//                    e.printStackTrace();
//                }
//            }
//
//            @Override
//            public void onError(Throwable error) {
//                log.error("发生连接错误：", error);  // 改成 log.error 打印堆栈
//                try {
//                    session.sendMessage(new TextMessage("[错误信息] " + error.getMessage()));
//                    session.close();
//                } catch (IOException e) {
//                    log.error("发送错误信息时失败：", e);
//                }
//            }
//        });
    }
}
